import os
from timm.models import hrnet
import os
from losses import new_lossfn
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import torch
from torch import optim
from model.edge_model import SegUNet
from dataset.fpv_1107_544 import Fpv
from torch.utils.data import DataLoader
from tqdm import tqdm
from torch import nn
import time
from utils.metrics import Metrics
from utils import clc
from collections import OrderedDict
import mlflow
mlflow.set_tracking_uri('http://192.168.10.78:5000')

def load_pretrain(check_file):
    od = torch.load(check_file)
    nd = OrderedDict()
    for k in od.keys():
        if 'mask_pred' not in k:
            nd[k] = od[k]
    return nd


def run():
    with mlflow.start_run(run_name='edge_model_1219_new_loss'):
        train_data = Fpv('/home/dsl/dataset/1219', is_train=True)
        test_data = Fpv('/home/dsl/dataset/1219', is_train=False)
        num_class = 9
        #train_data = Fpv(root='F:/dl_data/seg')
        trainloader = DataLoader(train_data, batch_size=8, shuffle=True, num_workers=8)
        testloader = DataLoader(test_data, batch_size=2, shuffle=False, num_workers=2)
        model = SegUNet(num_class)
        model.load_state_dict(load_pretrain('/data/check/xag_city_new286.pth'),strict=False)
        #model.load_state_dict(torch.load('/data/check/edge_model_1219_544att_144.pth'), strict=True)
        cre_fun1 = new_lossfn.SigCross(label_num=num_class)
        model.cuda()
        model.train()
        #optimizer = optim.AdamW(model.parameters(), lr=0.0002, weight_decay=2e-4)
        optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.9, weight_decay=2e-4)
        #optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=2e-4)
        #sch = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200, eta_min=2e-5)
        #sch = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 200, 500], gamma=0.1)
        sch = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=4, T_mult=2, eta_min=1e-6, verbose=True)
        #sch = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, min_lr=1e-6, patience=5)
        for epoch in range(600):
            losses = []
            sig_losses = []
            model.train()
            for img, target in tqdm(trainloader):
                img = torch.autograd.Variable(img.cuda())
                target = torch.autograd.Variable(target.cuda())
                pred = model(img)
                loss = cre_fun1(pred, target)
                #sig_loss = loss_f2(pred, target)
                total_loss = loss

                optimizer.zero_grad()
                total_loss.backward()

                #parameters = list(filter(lambda p: p.grad is not None, model.parameters()))
                #total_norm = max(p.grad.data.abs().max() for p in parameters)

                # print('total', total_norm, loss.item())
                nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0, norm_type=2)
                optimizer.step()
                losses.append(loss.item())
                #sig_losses.append(sig_loss.item())
            mlflow.log_metric('cross_loss', sum(losses)/len(losses), step=epoch)
            #mlflow.log_metric('sig_loss', sum(sig_losses) / len(sig_losses), step=epoch)
            sch.step()
            if epoch%2==0:
                try:
                    model.eval()
                    '''
                    me = Metrics(n_classes=num_class)
                    for img, target in tqdm(trainloader):
                        img = torch.autograd.Variable(img.cuda())
                        target = torch.autograd.Variable(target.cuda())
                        pred = model(img)
                        pred = torch.softmax(pred,dim=1)
                        pred = torch.argmax(pred, dim=1)
                        me.update(pred,target)
                    metr = me.compute_metrics()

                    print(metr)
                    mlflow.log_metric('train_miou', metr['miou'], step=epoch)
                    mlflow.log_metric('train_fw_miou', metr['fw_miou'], step=epoch)
                    '''
                    me = Metrics(n_classes=num_class)
                    for img, target in tqdm(testloader):
                        img = torch.autograd.Variable(img.cuda())
                        target = torch.autograd.Variable(target.cuda())
                        pred = model(img)
                        pred = torch.softmax(pred,dim=1)
                        pred = torch.argmax(pred, dim=1)
                        me.update(pred,target)
                    metr = me.compute_metrics()
                    print(metr)
                    mlflow.log_metric('val_miou', metr['miou'], step=epoch)
                    mlflow.log_metric('val_fw_miou', metr['fw_miou'], step=epoch)
                except Exception as e:
                    print(e)


            if epoch%2==0:
                torch.save(model.state_dict(), '/data/check/edge_model_1219_544_v3_' + str(epoch) + '.pth')




if __name__ == '__main__':
    run()
